package com.microframework.base.core.redis.service;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Service;
import lombok.extern.slf4j.Slf4j;
import java.util.Collections;

/**
 * Redis限流服务，基于令牌桶算法
 */
@Slf4j
@Service
public class RedisRateLimiterService {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    private static final String RATE_LIMITER_KEY_PREFIX = "rate_limiter:";

    // 固定窗口限流的Lua脚本
    private static final String FIXED_WINDOW_RATE_LIMITER_LUA = 
            "local key = KEYS[1] " +
            "local max_requests = tonumber(ARGV[1]) " +
            "local window_time = tonumber(ARGV[2]) " + 
            "local current = redis.call('get', key) " +
            "if current and tonumber(current) >= max_requests then " +
            "    return 0 " +  
            "end " +
            "local new_requests = 0 " +
            "if current then " + 
            "    new_requests = tonumber(current) + 1 " +
            "else " +
            "    new_requests = 1 " +
            "end " +
            "redis.call('set', key, new_requests, 'EX', window_time) " +
            "return 1";

    /**
     * 固定窗口限流
     *
     * @param key 限流标识
     * @param maxRequests 窗口期最大请求数
     * @param windowTimeSeconds 窗口期时间（秒）
     * @return 是否允许通过
     */
    public boolean tryAcquireFixedWindow(String key, int maxRequests, int windowTimeSeconds) {
        try {
            String redisKey = RATE_LIMITER_KEY_PREFIX + key;
            DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
            redisScript.setScriptText(FIXED_WINDOW_RATE_LIMITER_LUA);
            redisScript.setResultType(Long.class);
            
            Long result = redisTemplate.execute(
                    redisScript, 
                    Collections.singletonList(redisKey), 
                    maxRequests, 
                    windowTimeSeconds
            );
            
            boolean allowed = result != null && result == 1L;
            if (!allowed) {
                log.debug("限流生效 - key: {}, 最大请求数: {}, 窗口期: {}秒", key, maxRequests, windowTimeSeconds);
            }
            return allowed;
        } catch (Exception e) {
            log.error("限流器异常 - key: {}", key, e);
            // 出现异常时默认放行
            return true;
        }
    }
    
    // 漏桶算法限流的Lua脚本
    private static final String LEAKY_BUCKET_RATE_LIMITER_LUA = 
            "local key = KEYS[1]\n" +
            "local rate = tonumber(ARGV[1])\n" +
            "local capacity = tonumber(ARGV[2])\n" +
            "local now = tonumber(ARGV[3])\n" +
            "\n" +
            "local last_time = redis.call('hget', key, 'last_time')\n" +
            "local water_level = redis.call('hget', key, 'water_level')\n" +
            "\n" +
            "if last_time and water_level then\n" +
            "    last_time = tonumber(last_time)\n" +
            "    water_level = tonumber(water_level)\n" +
            "    local time_passed = math.max(0, now - last_time)\n" +
            "    local water_leaked = time_passed * rate\n" +
            "    water_level = math.max(0, water_level - water_leaked)\n" +
            "else\n" +
            "    water_level = 0\n" +
            "end\n" +
            "\n" +
            "local result = 0\n" +
            "if water_level < capacity then\n" +
            "    water_level = water_level + 1\n" +
            "    result = 1\n" +
            "end\n" +
            "\n" +
            "redis.call('hmset', key, 'last_time', now, 'water_level', water_level)\n" +
            "redis.call('expire', key, 60)\n" +
            "return result";

    /**
     * 漏桶算法限流
     *
     * @param key 限流标识
     * @param rate 漏桶流出速率（每秒处理多少请求）
     * @param capacity 漏桶容量
     * @return 是否允许通过
     */
    public boolean tryAcquireLeakyBucket(String key, int rate, int capacity) {
        try {
            String redisKey = RATE_LIMITER_KEY_PREFIX + "leaky:" + key;
            DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
            redisScript.setScriptText(LEAKY_BUCKET_RATE_LIMITER_LUA);
            redisScript.setResultType(Long.class);
            
            Long result = redisTemplate.execute(
                    redisScript, 
                    Collections.singletonList(redisKey), 
                    rate, 
                    capacity,
                    System.currentTimeMillis() / 1000
            );
            
            boolean allowed = result != null && result == 1L;
            if (!allowed) {
                log.debug("漏桶限流生效 - key: {}, 速率: {}/秒, 容量: {}", key, rate, capacity);
            }
            return allowed;
        } catch (Exception e) {
            log.error("漏桶限流器异常 - key: {}", key, e);
            // 出现异常时默认放行
            return true;
        }
    }
} 